package quantum2.graph;

import quantum2.math.Arithmetic;
import quantum2.math.Complex;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

/**
 * edge 0: A -> S (0 -> 2)
 * edge 1: B -> S (1 -> 2)
 * vertex 0 = A, vertex 1 = B, vertex 2 = S
 */
public final class Schedule {

    private static final class SmallEntry {

        final Complex amp;
        final int edge;
        final boolean forward;

        private SmallEntry(Complex amp, int edge, boolean forward) {
            this.amp = amp;
            this.edge = edge;
            this.forward = forward;
        }
    }

    private final Graph g;
    private final Arithmetic a;
    private final Matrix matrix;
    private final TreeMap<Number, List<SmallEntry>> list = new TreeMap<Number, List<SmallEntry>>();
    private Number currentTime;

    private final Object lock = new Object();

    public Schedule(Graph g, Arithmetic a, Matrix matrix) {
        this.g = g;
        this.a = a;
        this.matrix = matrix;
        this.currentTime = a.zero();
    }

    public void schedule(Number time, Complex amp, int edge, boolean forward) {
        synchronized (lock) {
            SmallEntry newEntry = new SmallEntry(amp, edge, forward);
            List<SmallEntry> existing = list.get(time);
            if (existing != null) {
                existing.add(newEntry);
            } else {
                Map.Entry<Number, List<SmallEntry>> upper = list.ceilingEntry(time);
                List<SmallEntry> addTo = null;
                if (upper != null) {
                    if (a.isNear(time, upper.getKey())) {
                        addTo = upper.getValue();
                    }
                }
                if (addTo == null) {
                    Map.Entry<Number, List<SmallEntry>> lower = list.floorEntry(time);
                    if (lower != null) {
                        if (a.isNear(time, lower.getKey())) {
                            addTo = lower.getValue();
                        }
                    }
                }
                if (addTo == null) {
                    List<SmallEntry> newList = new ArrayList<SmallEntry>();
                    newList.add(newEntry);
                    list.put(time, newList);
                } else {
                    addTo.add(newEntry);
                }
            }
        }
    }

    public void next() {
        synchronized (lock) {
            Number first = list.firstKey();
            List<SmallEntry> e = list.remove(first);
            currentTime = first.doubleValue();
            Complex ampToA = null;
            Complex ampToB = null;
            Complex ampAtoS = null;
            Complex ampBtoS = null;
            for (SmallEntry entry : e) {
                Complex amp = entry.amp;
                if (entry.edge == Graph.E_AS) {
                    if (entry.forward) {
                        ampAtoS = amp.accum(ampAtoS);
                    } else {
                        ampToA = amp.accum(ampToA);
                    }
                } else {
                    if (entry.forward) {
                        ampBtoS = amp.accum(ampBtoS);
                    } else {
                        ampToB = amp.accum(ampToB);
                    }
                }
            }
            {
                if (ampToA != null) {
                    Number len = g.getEdgeLength(Graph.E_AS);
                    Number nextTime = a.add(currentTime, len);
                    schedule(nextTime, ampToA.neg(), Graph.E_AS, true);
                }
            }
            {
                if (ampToB != null) {
                    Number len = g.getEdgeLength(Graph.E_BS);
                    Number nextTime = a.add(currentTime, len);
                    schedule(nextTime, ampToB.neg(), Graph.E_BS, true);
                }
            }
            {
                if (ampAtoS != null || ampBtoS != null) {
                    Complex ampStoA;
                    Complex ampStoB;
                    if (ampAtoS != null && ampBtoS != null) {
                        // case 2
                        ampStoA = matrix.K_FWD.mul(ampBtoS).add(matrix.K_AB.mul(ampAtoS));
                        ampStoB = matrix.K_FWD.mul(ampAtoS).add(matrix.K_BA.mul(ampBtoS));
                    } else if (ampAtoS != null) {
                        // case 1
                        ampStoA = matrix.K_AB.mul(ampAtoS);
                        ampStoB = matrix.K_FWD.mul(ampAtoS);
                    } else {
                        // case 3
                        ampStoA = matrix.K_FWD.mul(ampBtoS);
                        ampStoB = matrix.K_BA.mul(ampBtoS);
                    }
                    Number lenB = g.getEdgeLength(Graph.E_BS);
                    Number nextTimeB = a.add(currentTime, lenB);
                    Number lenA = g.getEdgeLength(Graph.E_AS);
                    Number nextTimeA = a.add(currentTime, lenA);
                    schedule(nextTimeB, ampStoB, Graph.E_BS, false);
                    schedule(nextTimeA, ampStoA, Graph.E_AS, false);
                }
            }
        }
    }

    public StatResult getStat() {
        double sum = 0;
        int count = 0;
        int nedges = g.getEdgeNum();
        double[] energy = new double[nedges];
        int[] edgeNum = new int[nedges];
        Number currentTime;
        synchronized (lock) {
            for (List<SmallEntry> entries : list.values()) {
                for (SmallEntry entry : entries) {
                    double e1 = entry.amp.conj();
                    int edge = entry.edge;
                    energy[edge] += e1;
                    edgeNum[edge]++;
                    sum += e1;
                    count++;
                }
            }
            currentTime = this.currentTime;
        }
        return new StatResult(currentTime, count, sum, energy, edgeNum, g);
    }

    public double getCurrentTime() {
        synchronized (lock) {
            return currentTime.doubleValue();
        }
    }
}
